require"curl"
require"lsqlite3"

SHOWS = {
	2930, -- Buffy The Vampire Slayer
	7926, -- Dexter
	15319, -- Californication
	12662, -- True Blood
	15614, -- Chuck
	8172,  -- Heroes
	18388, -- Fringe
	18967, -- The Mentalist
	15716, -- Pushing Daisies
	15619, -- The Secret Diary of a Call Girl
	15615, -- Life
	4696, -- Numb3rs
	4895, -- Prison Break
	18834, -- Merlin
	13445, -- Terminator-The Sarah Connor Chronicles
	2870, -- Bones
	15744, -- Reaper
	17816, -- Dollhouse
	4284, -- Lost
	15422, -- Damages
	2730, -- Battlestar Galactica
	10629, -- Kyle XY
	2445, -- 24
	15383, -- Burn Notice
	18531, -- Flashpoint
	7506, -- Eureka
	18751, -- In Plain Sight
	4432, -- Medium
	4677, -- Nip Tuck
	15313, -- Saving Grace
	8044, -- The IT Crowd
	18817, -- The Middlema
	11610, -- Life on Mars US
	6185, -- The Shield
	3332, -- Doctor Who 2005
	5047, -- Rome
	3267, -- Deadwood
	5410, -- Supernatural
	9165, -- Andy Barker P.I.
	15187, -- bionic woman
	10257, -- drive
	3683, -- gilmore girls
	14612, -- John from Cincinnati
	15189, -- journeyman
	6728, -- masters of horror
	8396, -- masters of sf
	18529, -- fear itself
	4964, --regenesis
	5514, -- The 4400
	5688, -- the collector
	6507, -- veronica mars
	15595, -- the kill point
	18933, -- eleventh hour
	18682, -- sanctuary
	2537, -- alias
}
local db = assert(sqlite3.open("shows.db"))
db:exec[[
CREATE TABLE IF NOT EXISTS shows (
	id INTEGER PRIMARY KEY,
	tvrage_id INTEGER,
	name TEXT NOT NULL
)
]]
db:exec[[
CREATE TABLE IF NOT EXISTS episodes (
	id INTEGER PRIMARY KEY,
	show_id INTEGER NOT NULL,
	episode_number INTEGER,
	season_number INTEGER,
	episode_number_in_season INTEGER,
	production_id INTEGER,
	title TEXT NOT NULL,
	air_date DATETIME
)]]

local function get_basic_curl(add_header)
	local c = curl.easy_init()
	local output = {}
	c:setopt(curl.OPT_WRITEFUNCTION, function (s)
		output[#output + 1] = s
		return s:len()
	end)
	if add_header then
		local header = {}
		c:setopt(curl.OPT_HEADERFUNCTION, function (s)
			header[#header + 1] = s
			return s:len()
		end)
		return c, output, header
	else
		return c, output
	end
end

local function get_response_code(header)
	local version, code, message = header:match("(HTTP/%d+\.%d+)%s+([0-9]+) (.*)\n")
	code = tonumber(code)
	return code, message, version
end

local function get(url)
	local c, output, header = get_basic_curl(true)
	c:setopt(curl.OPT_URL, url)
	assert(c:perform())
	local code, message, version = get_response_code(header[1])
	assert(code < 400, "HTTP fail")
	return table.concat(output)
end

local function get_show_name(s)
	local p = s:find("<tr bgcolor='#521923' >", 1, true)
	local start = s:find("<b>", p, true) + 3
	local stop = s:find("</b>", start, true) - 1
	return s:sub(start, stop)
end

local function fix_air_date(entry)
	local start = entry:find("<table align='center'", nil, true)
	if not start then return entry end
	local stop = entry:find("</table>", start, true) + 8
	assert(stop)
	local date_values = entry:sub(start, stop)
	local values = {}
	for entry in date_values:gmatch("<td[^>]*>(.-)</td>") do
		values[#values + 1] = entry
	end
	return entry:sub(1, start - 1)..table.concat(values)..entry:sub(stop)
end

local function remove_link(entry)
	entry = entry:gsub("^%s+", "")
	if entry:sub(1, 2) ~= "<a" then return entry end
	local start = entry:find(">", 1, true) + 1
	local stop = entry:find("</a>", 1, true) - 1
	return entry:sub(start, stop)
end

local function remove_img(entry)
	local start = entry:find("<img", 1, true)
	if not start then return entry end
	local stop = entry:find("</a>", start, true) + 4
	return entry:sub(stop)
end

local get_show_id_stmt = db:prepare"SELECT id FROM shows WHERE name = ?"
local insert_show_stmt = db:prepare"INSERT INTO shows (name, tvrage_id) VALUES (?, ?)"
local function get_show_id(name, tvrage_id)
	get_show_id_stmt:bind_values(name)
	local result = get_show_id_stmt:step()
	if result == sqlite3.ROW then
		result = get_show_id_stmt:get_value(0)
		get_show_id_stmt:reset()
		return tonumber(result)
	else
		get_show_id_stmt:reset()
	end
	insert_show_stmt:bind_values(name, tvrage_id)
	insert_show_stmt:step()
	insert_show_stmt:reset()
	return db:last_insert_rowid()
end

local get_show_id_by_rageid_stmt = db:prepare"SELECT id FROM shows WHERE tvrage_id = ?"
local function get_show_id_by_tvrage_id(tvrage_id)
	get_show_id_by_rageid_stmt:bind_values(tvrage_id)
	local result = get_show_id_stmt:step()
	if result == sqlite3.ROW then
		result = get_show_id_by_rageid_stmt:get_value(0)
		get_show_id_by_rageid_stmt:reset()
		return tonumber(result)
	else
		get_show_id_by_rageid_stmt:reset()
	end
end

local clear_show_stmt = assert(db:prepare"DELETE FROM episodes WHERE show_id = ?", db:errmsg())
local insert_episode_stmt = assert(db:prepare[[
INSERT INTO episodes (
	show_id,
	season_number,
	episode_number,
	episode_number_in_season,
	production_id,
	air_date,
	title
) VALUES (?, ?, ?, ?, ?, ?, ?)
]], db:errmsg())
local function get_episode_list(s, id)
	clear_show_stmt:bind_values(id)
	clear_show_stmt:step()
	clear_show_stmt:reset()

	local hr_pos = s:find("<hr style='border: 2px solid rgb(102, 102, 102); margin-top: 1px; margin-bottom: 1px;'>", nil, true)
	local start_season = s:find("<font size='4'>", hr_pos, true)
	while start_season do
		local end_season = s:find("<font size='4'>", start_season + 1, true)
		local end_season_name = s:find("</font>", start_season, true)
		local season_name = s:sub(start_season + 15, end_season_name - 1)

		local season_number = season_name:match("Season (%d+)")

		local start_episode = s:find("<tr bgcolor='[^']+'  id=\"brow\"", start_season + 1)
		while start_episode and start_episode < (end_season or s:len()) do
			local end_episode = s:find("<tr bgcolor='[^']+'  id=\"brow\"", start_episode + 1)
			local episode_entry = s:sub(start_episode + 13, end_episode)
			episode_entry = fix_air_date(episode_entry)
			local data, temp = {}, {}
			for entry in episode_entry:gmatch("<td[^>]*>(.-)</td>") do
				temp[#temp + 1] = entry
			end
			if #temp < 5 then break end
			data[1] = tonumber(remove_link(temp[1])) -- episode number
			data[2] = tonumber(remove_link(temp[2]):match("%d+x(%d+)")) -- episode number in season
			data[3] = temp[3] -- prod number
			data[4] = temp[4] and os.date(temp[4]) -- air date
			data[5] = remove_link(remove_img(temp[5])) -- episode title

			insert_episode_stmt:bind_values(id, season_number, unpack(data))
			insert_episode_stmt:step()
			insert_episode_stmt:reset()

			start_episode = end_episode
		end

		start_season = end_season
	end
end

for _, show in ipairs(SHOWS) do
	assert(type(show) == "number")
	if not get_show_id_by_tvrage_id(show) then
		local url = string.format("http://www.tvrage.com/shows/id-%d/episode_list/all", show)
		local s = get(url)
		local name = get_show_name(s)
		local show_id = get_show_id(name, show)
		print(name) io.flush()
		get_episode_list(s, show_id)
	end
end
